from typing import List
import numpy as np  
import torch
import torch.nn.functional as F
from PIL import Image  
from utils import * 
from torchvision import transforms
from UI import plot_sudoku


# TODO:各种数独图像的识别兼容性


def Identify_number(img_path: str):
    '''
    input: path of imgae
    output: sudu number
    '''
    
    # 初始化棋盘
    board = np.zeros((9, 9), dtype=int)
    
    img    = Image.open(img_path)
    box_81, tag_81 = get_box_81(img_path)
    
    # 确定每行每列一个标准位置
    row_index    = np.array(sorted(box_81, key=lambda x: x[0]))[:,0].reshape(9,9).mean(axis=1)
    row_index    = np.round(row_index, 0)
    column_index = np.array(sorted(box_81, key=lambda x: x[1]))[:,1].reshape(9,9).mean(axis=1)
    column_index = np.round(column_index, 0)
    
    for box, tag in zip(box_81, tag_81):
        if tag == 0:
            continue
        else:
            img_box = img.crop(box) #裁剪出单元格图像
            row, column = location(box, row_index, column_index)
            board[row][column]   = number_identify(img_box)

    return board

def number_identify(img : Image.Image):
    '''
    input : PIL格式图像
    output: 识别的数字
    '''
    
    def check_background(img):
        '''
        确保图片是黑底白字
        '''
        img = img.convert('L')  
        img_copy = img.resize((10, 10))
        img_array = np.array(img_copy)
        mid_threshold = 120   
        black_pixels = np.sum(img_array < mid_threshold)
        white_pixels = np.sum(img_array >= mid_threshold)
        if black_pixels > white_pixels:
            return img
        else:
            return img.point(lambda x: 255 - x)
    
    def process_image(img : Image.Image):
        img = img.convert("RGB")
        transform = transforms.Compose([transforms.Resize((28,28)),
                                                transforms.Grayscale(),
                                                transforms.ToTensor(),
                                                transforms.Normalize((0.5,),(0.5,))])
        img_tensor = transform(img)
        img_tensor = img_tensor.unsqueeze(0)  # 添加 batch 维度
        return img_tensor

    img = check_background(img)
    img = process_image(img)
    with torch.no_grad():
        output = model(img)
        predicted_class = torch.argmax(F.softmax(output[0], dim=0)).item()
        return predicted_class
    
# 数独求解算法：回溯法
class Solution:
    def solveSudoku(self, board: List[List[str]]) -> None:
        self.backtrack(board, 0, 0)
    
    def backtrack(self, board, i, j):
        m, n = 9, 9
        if j == n:
            # 遍历到最后一列就换到下一行
            return self.backtrack(board, i +1, 0)
    
        if i== m:
            return True
        
        # 碰到已经填好的数字直接跳过
        if board[i][j] != 0:
            return self.backtrack(board, i, j + 1)

        for number in range(1, 10):
            # 数字不合法就跳过
            if not self.isValid(board, i, j, number):
                continue
            
            board[i][j] = number
            
            # 如果找到一个解就结束
            if self.backtrack(board, i, j+1):
                return True
            # 如果没找到，重新变成空白格
            board[i][j] = 0
        
        return False
        
    # 数字合法性判断
    def isValid(self, board, row, col, number):
        for i in range(9):
            if board[row][i] == number:
                return False
            if board[i][col] == number:
                return False
            if board[(row // 3) * 3 + i // 3][(col // 3) * 3 + i % 3] == number:
                return False
        return True


    
if __name__ == '__main__':


    model = torch.load('LeNet.pth', map_location=torch.device('cpu'))   
    board = Identify_number('image/2.jpg')
    board_init = board.copy()  
    

    Solution().solveSudoku(board)
    plot_sudoku(board, board_init)


            

    